package network

import (
	"net"
	"operator/pkg/codes/tcp"
	"sync"
)

type TcpClient struct {
	conn         net.Conn
	responseDict map[string]chan []byte
	mu           sync.RWMutex
	dictFn       map[string]func([]byte) []byte
}

func (tc *TcpClient) List() ([]byte, error) {
	data := tcp.Encode("list", nil)
	id := data.ID()
	_, err := tc.conn.Write(data.ToBytes())
	if err != nil {
		return nil, err
	}

	ch := make(chan []byte)
	tc.mu.Lock()
	tc.responseDict[id] = ch
	tc.mu.Unlock()

	result := <-ch

	tc.mu.Lock()
	delete(tc.responseDict, id)
	tc.mu.Unlock()

	return result, nil
}

func (tc *TcpClient) Write(tag string, data []byte) ([]byte, error) {
	dd := tcp.Encode(tag, data)
	id := dd.ID()
	_, err := tc.conn.Write(dd.ToBytes())
	if err != nil {
		return nil, err
	}

	ch := make(chan []byte)
	tc.mu.Lock()
	tc.responseDict[id] = ch
	tc.mu.Unlock()

	result := <-ch

	return result, nil
}

func (tc *TcpClient) read() error {
	buf := make([]byte, 1024*30)
	n, err := tc.conn.Read(buf)
	if err != nil {
		return err
	}
	tdata := tcp.Decode(buf[:n])
	if tdata.Tag().IsResponse() {
		tc.mu.RLock()
		tc.responseDict[tdata.ID()] <- tdata.Data()
		tc.mu.RUnlock()
	} else {
		if tdata.Tag().String() == "probe" {
			d := tdata.ResponseData([]byte{1})
			tc.conn.Write(d.ToBytes())
		} else {
			if fn, ok := tc.dictFn[tdata.Tag().String()]; ok {
				res := fn(tdata.Data())
				resdata := tdata.ResponseData(res)
				_, err = tc.conn.Write(resdata.ToBytes())
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func (ts *TcpClient) AddFn(tag string, fn func([]byte) []byte) {
	ts.dictFn[tag] = fn
}

func (tc *TcpClient) Run(stopCh <-chan struct{}) {
	for {
		select {
		case <-stopCh:
			break
		default:
			err := tc.read()
			if err != nil {
				return
			}
		}
	}
}

func NewTcpClient(addr string) (*TcpClient, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	tc := &TcpClient{
		conn:         conn,
		responseDict: map[string]chan []byte{},
		mu:           sync.RWMutex{},
		dictFn:       map[string]func([]byte) []byte{},
	}

	return tc, nil
}
